import gymnasium as gym
from gymnasium import spaces
import numpy as np
import mujoco
import glfw

class YmbotEPIDEnv(gym.Env):
    def __init__(self, render_mode=None, window_width=1200, window_height=900):
        super(YmbotEPIDEnv, self).__init__()
        # 加载 ymbot_e 的 Mujoco 模型
        self.model = mujoco.MjModel.from_xml_path('/home/sh/catkin_ws/src/ymbot_e_description/mjcf/ymbot_e_description.xml')
        self.data = mujoco.MjData(self.model)
        self.render_mode = render_mode
        # 定义自定义环境的动作空间，90 个 PID 参数（每个关节3个参数），每个参数范围为 [0, 10]
        self.action_space = spaces.Box(low=0, high=10, shape=(90,), dtype=np.float32)
        # 定义自定义环境的观测空间，30 个关节的角度，每个角度范围为 [-10, 10]
        self.observation_space = spaces.Box(low=-10, high=10, shape=(30,), dtype=np.float32)
        # 目标值曲线
        self.setpoint_curve = lambda t: np.zeros(30)
        # 当前状态
        self.state = np.zeros(30)
        # 上一次的误差
        self.last_error = np.zeros(30)
        # 积分项
        self.integral = np.zeros(30)
        # 时间步
        self.timestep = 0
        self.control_output = np.zeros(30)  # 初始化控制输出
        self.mouse_button_left = False
        self.mouse_button_middle = False
        self.mouse_button_right = False
        self.last_mouse_x = 0
        self.last_mouse_y = 0
        self.scene = None
        self.camera = mujoco.MjvCamera()
        self.window_width = window_width
        self.window_height = window_height

    def set_target_curve(self, curve_function):
        """设置目标曲线"""
        self.setpoint_curve = curve_function

    def reset(self, seed=None, options=None):
        # 重置 ymbot_e 环境
        mujoco.mj_resetData(self.model, self.data)
        # 随机初始化当前状态
        self.state = np.random.uniform(-1, 1, size=30)
        # 重置误差和积分项
        self.last_error = np.zeros(30)
        self.integral = np.zeros(30)
        # 重置时间步
        self.timestep = 0
        return self.state, {}

    def step(self, action):
        # 提取 PID 参数
        kp = action[:30]
        ki = action[30:60]
        kd = action[60:90]
        
        # 更新目标值
        self.setpoint = self.setpoint_curve(self.timestep)
        
        # 计算误差、积分和导数项
        error = self.setpoint - self.state
        self.integral += error
        derivative = error - self.last_error
        self.last_error = error
        # 计算控制量
        self.control_output = kp * error + ki * self.integral + kd * derivative
        
        # 将控制量应用到 ymbot_e 环境的所有关节
        self.data.ctrl[:] = np.clip(self.control_output, -1, 1)
        mujoco.mj_step(self.model, self.data)
        
        # 更新状态
        self.state = self.data.qpos[:30]  # 使用观测空间中的前30个元素作为关节角度
        # 奖励为负的绝对误差之和
        reward = -np.sum(np.abs(error))
        # 当所有误差小于阈值时，认为任务完成
        done = np.all(np.abs(error) < 0.01)
        
        # 增加时间步
        self.timestep += 1
        
        return self.state, reward, done, False, {}

    def get_control_output(self):
        """获取控制输出"""
        return self.control_output

    def mouse_button_callback(self, window, button, action, mods):
        if button == glfw.MOUSE_BUTTON_LEFT:
            self.mouse_button_left = (action == glfw.PRESS)
        elif button == glfw.MOUSE_BUTTON_MIDDLE:
            self.mouse_button_middle = (action == glfw.PRESS)
        elif button == glfw.MOUSE_BUTTON_RIGHT:
            self.mouse_button_right = (action == glfw.PRESS)
        self.last_mouse_x, self.last_mouse_y = glfw.get_cursor_pos(window)

    def mouse_move_callback(self, window, xpos, ypos):
        if not self.mouse_button_left and not self.mouse_button_middle and not self.mouse_button_right:
            return

        dx = xpos - self.last_mouse_x
        dy = ypos - self.last_mouse_y
        self.last_mouse_x = xpos
        self.last_mouse_y = ypos

        width, height = glfw.get_window_size(window)
        mod_shift = (glfw.get_key(window, glfw.KEY_LEFT_SHIFT) == glfw.PRESS or
                     glfw.get_key(window, glfw.KEY_RIGHT_SHIFT) == glfw.PRESS)

        if self.mouse_button_right:
            action = mujoco.mjtMouse.mjMOUSE_MOVE_H if mod_shift else mujoco.mjtMouse.mjMOUSE_MOVE_V
        elif self.mouse_button_left:
            action = mujoco.mjtMouse.mjMOUSE_ROTATE_H if mod_shift else mujoco.mjtMouse.mjMOUSE_ROTATE_V
        else:
            action = mujoco.mjtMouse.mjMOUSE_ZOOM

        mujoco.mjv_moveCamera(self.model, action, dx / height, dy / height, self.scene, self.camera)

    def scroll_callback(self, window, xoffset, yoffset):
        mujoco.mjv_moveCamera(self.model, mujoco.mjtMouse.mjMOUSE_ZOOM, 0, -0.05 * yoffset, self.scene, self.camera)

    def render(self):
        if self.render_mode == 'human':
            if not hasattr(self, 'context'):
                if not glfw.init():
                    raise Exception("GLFW can't be initialized")
                self.window = glfw.create_window(self.window_width, self.window_height, "Ymbot E PID Env", None, None)
                if not self.window:
                    glfw.terminate()
                    raise Exception("GLFW window can't be created")
                glfw.make_context_current(self.window)
                self.context = mujoco.MjrContext(self.model, mujoco.mjtFontScale.mjFONTSCALE_150)
                glfw.set_mouse_button_callback(self.window, self.mouse_button_callback)
                glfw.set_cursor_pos_callback(self.window, self.mouse_move_callback)
                glfw.set_scroll_callback(self.window, self.scroll_callback)
                self.scene = mujoco.MjvScene(self.model, maxgeom=1000)
            viewport = mujoco.MjrRect(0, 0, self.window_width, self.window_height)
            mujoco.mjv_updateScene(self.model, self.data, mujoco.MjvOption(), None, self.camera, mujoco.mjtCatBit.mjCAT_ALL, self.scene)
            mujoco.mjr_render(viewport, self.scene, self.context)
            glfw.swap_buffers(self.window)
            glfw.poll_events()

    def close(self):
        if hasattr(self, 'window'):
            glfw.destroy_window(self.window)
            glfw.terminate()

if __name__ == "__main__":
    env = YmbotEPIDEnv(render_mode='human', window_width=2560, window_height=1440)
    state, _ = env.reset()
    for _ in range(1000):
        action = np.random.uniform(low=env.action_space.low, high=env.action_space.high)  # 生成随机动作
        state, reward, done, truncated, info = env.step(action)
        # print(f"Action: {action}")  # 打印动作
        # print(f"State: {state}")  # 打印状态
        print(f"Reward: {reward}")  # 打印奖励
        env.render()
        if done:
            state, _ = env.reset()  # 重新初始化环境状态
    env.close()
